
#include "div_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>


namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t TILE_NUM = 8;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  DivTilingData tiling;
    
  int32_t NUM = 10;
  auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  uint64_t ubSize;
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
  auto aivNum = ascendcPlatform.GetCoreNum();

  uint32_t totalLength = 0, minLength = context->GetInputTensor(0)->GetShapeSize();
  for (int i = 0; i < 3; ++i) {
      totalLength = std::max<uint32_t>(totalLength, context->GetInputTensor(i)->GetShapeSize());
      minLength = std::min<uint32_t>(minLength, context->GetInputTensor(i)->GetShapeSize());
  }

    uint32_t x1Length = context->GetInputTensor(0)->GetShapeSize();
    uint32_t x2Length = context->GetInputTensor(1)->GetShapeSize();
    auto dt = context->GetInputTensor(0)->GetDataType();
    uint32_t sizeofdatatype;
    if (dt == ge::DT_INT8) {
        sizeofdatatype = 1;
        NUM = 12;
    }
    else if (dt == ge::DT_FLOAT16 || dt == ge::DT_BF16) {
        sizeofdatatype = 2;
    }
    else {
        sizeofdatatype = 4;
    }

    uint32_t alignNum = BLOCK_SIZE / sizeofdatatype;
    uint32_t tiling_size = ((ubSize) / BLOCK_SIZE / 2) / NUM;
    tiling_size = tiling_size <= 8 ? tiling_size : tiling_size / 8 * 8;

    uint32_t blockSize = tiling_size * alignNum;
    if (totalLength != minLength) {
        blockSize = std::min(blockSize, minLength);
        while (minLength % blockSize || minLength % alignNum) {
            blockSize -= 1;
        }
    }

    aivNum = (aivNum < totalLength / blockSize) ? aivNum : (totalLength / blockSize);
    aivNum = aivNum >= 1 ? aivNum : 1;

    uint32_t coreSize = (totalLength / aivNum) / (alignNum * 8) * (alignNum * 8);
    uint32_t coreRemain = totalLength - aivNum * coreSize;

    tiling.set_alignNum(alignNum);
    tiling.set_blockSize(blockSize);
    tiling.set_aivNum(aivNum);
    tiling.set_coreSize(coreSize);
    tiling.set_coreRemain(coreRemain);
    tiling.set_totalLength(totalLength);
    tiling.set_x1Length(x1Length);
    tiling.set_x2Length(x2Length);

    context->SetBlockDim(aivNum);

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Div : public OpDef {
public:
    explicit Div(const char* name) : OpDef(name)
    {
        this->Input("x1")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Input("x2")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_INT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b")
        .AddConfig("ascend310b");

    }
};

OP_ADD(Div);
}
